import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics import classification_report

# 加载数据
train_data = pd.read_csv('ChatGPT生成文本检测器公开数据-更新/train.csv')
test_data = pd.read_csv('ChatGPT生成文本检测器公开数据-更新/test.csv')

# 数据预处理
train_data['content'] = train_data['content'].apply(lambda x: x[1:-1])
test_data['content'] = test_data['content'].apply(lambda x: x[1:-1])

# 划分训练集和验证集
train_text, valid_text, train_label, valid_label = train_test_split(
    train_data['content'], train_data['label'], test_size=0.2, random_state=42
)

# TF-IDF向量化
tfidf = TfidfVectorizer(token_pattern=r'\w{1,}', max_features=5000, ngram_range=(1, 2))
train_tfidf = tfidf.fit_transform(train_text)
valid_tfidf = tfidf.transform(valid_text)
test_tfidf = tfidf.transform(test_data['content'])

# 模型训练和评估
model = LogisticRegression(max_iter=1000)  # 可以调整更多参数
model.fit(train_tfidf, train_label)

train_predictions = model.predict(train_tfidf)
valid_predictions = model.predict(valid_tfidf)

# 打印训练集和验证集的分类报告
print("Training Set Classification Report:")
print(classification_report(train_label, train_predictions))

print("Validation Set Classification Report:")
print(classification_report(valid_label, valid_predictions))

# 使用模型进行测试集预测
test_predictions = model.predict(test_tfidf)
test_data['label'] = test_predictions

# 保存预测结果
test_data[['name', 'label']].to_csv('tfidf_predictions.csv', index=None)
